import pickle
import cv2
import numpy as np
import os
import json

from utils import y16formatvisionutil

# import y16formatvisionutil


class Pose:
    """
    用于加载和处理猪的体尺数据
    """

    def __init__(
        self,
        pig_pose: list[dict],
        img_wh: list[float, float],
        pig_xywh: list[float, float, float, float],
        weight: float,
        depth_img: np.array,
        pig_mask_points:list[list]
    ):
        self.l2 = []
        self.m = []
        self.n = []
        self.l1 = []
        self.s1 = []
        self.s2 = []
        self.h1 = []
        self.h2 = []
        self.o = []
        self.depth = depth_img
        self.pig_mask_points = pig_mask_points
        if len(set([i.get("group_id", 0) for i in pig_pose])) < 9:
            raise Exception("points error")
        for d in pig_pose:
            group_id = d["group_id"]
            if group_id is None:
                group_id = 0
            points = d["points"][0]
            if len(points) == 0:
                raise Exception("points error")
            if group_id == 0:
                self.l2 = points
            elif group_id == 1:
                self.m = points
            elif group_id == 2:
                self.n = points
            elif group_id == 3:
                self.l1 = points
            elif group_id == 4:
                self.s1 = points
            elif group_id == 5:
                self.s2 = points
            elif group_id == 6:
                self.h1 = points
            elif group_id == 7:
                self.h2 = points
            elif group_id == 8:
                self.o = points

        [self.img_width, self.img_height] = img_wh
        self.weight = weight
        [self.pig_x, self.pig_y, self.pig_w, self.pig_h] = pig_xywh

    def body_size(self, scale: float = None) -> list[float]:
        """
        读取体尺数据

        Args:
            scale: 缩放比例
        """
        if scale is None:
            scale = 1.0
        return [
            # 体尺点距离数据
            self.__dist(self.l2, self.m, scale),
            self.__dist(self.m, self.o, scale),
            self.__dist(self.o, self.n, scale),
            self.__dist(self.n, self.l1, scale),
            self.__dist(self.s1, self.m, scale),
            self.__dist(self.m, self.s2, scale),
            self.__dist(self.h1, self.n, scale),
            self.__dist(self.n, self.h2, scale),
            # 体尺点关键角度数据
            self.__angle(self.l2, self.m, self.s1, scale),
            self.__angle(self.l2, self.m, self.s2, scale),
            self.__angle(self.o, self.m, self.s1, scale),
            self.__angle(self.o, self.m, self.s2, scale),
            self.__angle(self.m, self.o, self.n, scale),
            self.__angle(self.o, self.n, self.h1, scale),
            self.__angle(self.o, self.n, self.h2, scale),
            self.__angle(self.l1, self.n, self.h1, scale),
            self.__angle(self.l1, self.n, self.h2, scale),
            # 图片大小
            self.img_width * scale,
            self.img_height * scale,
            # 猪的位置
            self.pig_x * scale,
            self.pig_y * scale,
            self.pig_w * scale,
            self.pig_h * scale,
        ]

    def depth_img(self, scale: float) -> np.array:
        # 通过蒙版, 只提取猪的深度数据
        pig_area_points = np.array(self.pig_mask_points, dtype=np.int32)
        mask_img = np.zeros((self.img_height, self.img_width), dtype=np.uint8)
        mask = cv2.fillPoly(mask_img, [pig_area_points], 1)
        pig_only_depth = mask * self.depth

        dsize = (int(self.img_width * scale), int(self.img_height * scale))
        resize = cv2.resize(pig_only_depth, dsize)
        depth = resize.reshape((dsize[1], dsize[0], 1))
        return depth

    def __angle(
        self,
        a: list[float, float],
        o: list[float, float],
        b: list[float, float],
        scale=None,
    ):
        """
        计算∠aob的角度

        Args:
          a: a点的坐标
          o: o点的坐标
          b: b点的坐标
        Returns:
          ∠aob的角度, 弧度制
        """
        if scale is None:
            scale = 1.0
        A = np.array(a) * scale
        O = np.array(o) * scale
        B = np.array(b) * scale
        OA = A - O
        OB = B - O
        dot_product = np.dot(OA, OB)
        magnitude_OA = np.linalg.norm(OA)
        magnitude_OC = np.linalg.norm(OB)
        cos_theta = dot_product / (magnitude_OA * magnitude_OC)
        angle = np.arccos(cos_theta)
        return angle

    def __dist(self, point1, point2, scale=None):
        """
        计算两点的欧几里得距离

        Args:
          point1: 坐标
          point2: 坐标

        Returns:
          两点的欧几里得距离
        """
        # 将点转换为 NumPy 数组
        if scale is None:
            scale = 1.0
        p1 = np.array(point1) * scale
        p2 = np.array(point2) * scale

        # 计算欧几里得距离
        distance = np.linalg.norm(p1 - p2)
        return distance

    def load_data(data_root: str) -> list:
        # 遍历目录及其子文件夹
        data_list = []
        weight_pre = 0
        for root, dirs, files in os.walk(data_root):
            for filename in files:
                # 检查文件是否是 JSON 文件
                if not filename.endswith(".json"):
                    continue
                pig_pose = []
                img_wh = []
                pig_mask = []
                weight = 0.0
                pig_pose, img_wh, pig_mask = Pose.__load_json(root, filename)
                weight = Pose.__load_meta(root)
                if weight == 0:
                    weight = weight_pre
                weight_pre = weight
                pig_depth_img = Pose.__load_depth(root, img_wh)
                if len(pig_pose) == 9:
                    pig_mask_points = pig_mask["points"]
                    x, y, w, h = cv2.boundingRect(
                        np.array(pig_mask_points, dtype=np.float32)
                    )
                    data_list.append(
                        Pose(
                            pig_pose,
                            img_wh,
                            [x, y, w, h],
                            weight,
                            pig_depth_img,
                            pig_mask_points,
                        )
                    )
        return data_list

    def __load_json(root, filename):
        json_file_path = os.path.join(root, filename)
        # 读取 JSON 文件
        with open(json_file_path, "r", encoding="utf-8") as json_file:
            json_data = json.load(json_file)
            img_wh = [json_data["imageWidth"], json_data["imageHeight"]]
            shapes = json_data["shapes"]
            pig_pose = [i for i in shapes if i["label"] == "pig_pose"]
            pig_mask = [i for i in shapes if i["label"] == "pig"][0]

        return pig_pose, img_wh, pig_mask

    def __load_meta(root):
        # 构建 META_INFO.txt 文件路径
        meta_info_path = os.path.join(root, "META-INFO.txt")
        # 读取 META_INFO.txt 文件
        with open(meta_info_path, "r", encoding="utf-8") as meta_file:
            meta_info = meta_file.read()
            return float(meta_info.split(",")[4].replace("kg", ""))

    def __load_depth(root, img_wh):
        # 读取深度数据
        depth_info_path = os.path.join(root, "1280x720_MA.raw")
        depth_img = y16formatvisionutil.read_raw_depth_image(
            depth_info_path, img_wh[0], img_wh[1]
        )
        return depth_img

    pass


if "__main__" == __name__:
    data_list = Pose.load_data(f"C:/Users/kang_/Desktop/data")
    # 将对象序列化并写入文件
    with open("data.pkl", "wb") as file:
        pickle.dump(data_list, file)
